package org.example;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.event.EventListener;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.ExecutorSubscribableChannel;
import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.PingMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.*;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
import org.springframework.web.socket.messaging.SessionConnectedEvent;
import org.springframework.web.socket.messaging.SessionDisconnectEvent;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

import java.io.IOException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;


/**
 * 此 config 可以正常运行，单机情况下，测试了1晚上，吴断链现场
 */
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig2 implements WebSocketMessageBrokerConfigurer {

    // STOMP 心跳线程池
    @Bean(destroyMethod = "shutdown")
    public ThreadPoolTaskScheduler stompHeartbeatScheduler() {
        ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
        scheduler.setThreadNamePrefix("stomp-hb-");
        scheduler.initialize();
        return scheduler;
    }

    // WebSocket Ping线程池
    @Bean(destroyMethod = "shutdownNow")
    public ScheduledExecutorService webSocketPingScheduler() {
        return Executors.newScheduledThreadPool(1,
                new CustomizableThreadFactory("ws-ping-"));
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        registry.enableSimpleBroker("/topic")
                .setTaskScheduler(stompHeartbeatScheduler())
                //[server->client,client->server]
                .setHeartbeatValue(new long[]{0, 20000});
        registry.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws").setAllowedOriginPatterns("*");
    }

    @Bean
    public WebSocketHandlerDecoratorFactory pingDecoratorFactory(
            ScheduledExecutorService webSocketPingScheduler) {

        return handler -> new WebSocketHandlerDecorator(handler) {
            @Override
            public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                webSocketPingScheduler.scheduleAtFixedRate(() -> {
                    if (session.isOpen()) {
                        try {
                            session.sendMessage(new PingMessage());
                        } catch (IOException e) {
                            System.err.println("Ping error: " + e.getMessage());
                        }
                    }
                }, 10, 10, TimeUnit.SECONDS);

                super.afterConnectionEstablished(session);
            }
        };
    }

    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registration) {
        registration.addDecoratorFactory(pingDecoratorFactory(webSocketPingScheduler()))
                .setSendTimeLimit(15000);
    }

    @Bean
    public ServletServerContainerFactoryBean webSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        //        maxSessionIdleTimeout的核心作用
        //定义：设置WebSocket会话在完全无数据交互情况下的最大存活时间（单位：毫秒）。
        //示例值30000L表示：如果连接30秒内无任何数据（包括Ping/Pong、业务消息)，服务器会主动断开连接。
        container.setMaxSessionIdleTimeout(60*1000L); // 1分钟
        return container;
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);

                // 心跳消息的识别标准：
                // 1. 无STOMP命令
                // 2. 消息体为空
                // 3. 包含heartbeat头
                if (accessor.getCommand() == null &&
                        message.getPayload() instanceof byte[] &&
                        ((byte[]) message.getPayload()).length == 0) {
                    System.out.println("[STOMP心跳] 类型: " +
                            (channel instanceof ExecutorSubscribableChannel ? "接收" : "发送"));
                }
                return message;
            }
        });
    }

}